#! /usr/bin/env python
# coding=utf-8

# 将 aa比对 映射到 对应的nt中
# 注意必须在alicut之前运行这一步 不然就把 所有的 ALICUT* 文件 move到一个地方

# nucl 和 aa的文件名称必须一样 放在两个不同的路径下

import os
import threading
import shutil
import sys
import argparse
from Bio import SeqIO



parser = argparse.ArgumentParser(
    description='''将核酸序列 根据其氨基酸的比对 按 codon进行排序
    用法:
    align_aa2nt.py -a Orthogroup_Sequences_mafft_raw -n Orthogroup_Sequences_namechanged_nucl -o Orthogroup_Sequences_mafft_nucl
    由大天才于2021年7月20日创建于浙江农业大学''')

parser.add_argument('-a',
                help='必须给定，比对好的氨基酸的路径')


parser.add_argument('-o',
                help='必须给定，输出文件的路径')

parser.add_argument('-n',
                help='必须给定，没比对的核酸的路径')



args = parser.parse_args()

if not args.a or not args.o or not args.n:
    parser.print_help()
    sys.exit()



infile_aa = args.a

outfile_nucl  = args.o

infile_nucl = args.n






try:
	shutil.rmtree(outfile_nucl)
except:
	pass

try:
	os.mkdir(outfile_nucl)
except:
	pass






for i in os.listdir(infile_aa):
	if i .split('.')[-1] in ['fa','fasta']:
		# 读取对应的核酸序列
		nucl_dic = {}

		outfile = open(outfile_nucl+'/'+i,'w')

		for j in SeqIO.parse(infile_nucl+'/'+i,'fasta'):
			nucl_dic[str(j.name)] = str(j.seq)

		# 读取比对序列
		for j in SeqIO.parse(infile_aa+'/'+i,'fasta'):
			matched_nucl_seq = nucl_dic[str(j.name)]
			new_seq = ''
			start_point = 0
			for k in j.seq:
				if k =='-':
					new_seq+='---'
				else:
					new_seq += matched_nucl_seq[start_point:start_point+3]
					start_point += 3

			outfile.write('>'+j.name+'\n')
			outfile.write(new_seq+'\n')

		outfile.close()



